#  ICE Revision: $Id$
"""Get additional data

This module gets additional data from an OpenFOAM-case and generates a vtkTable-object
"""

from PyFoam.Paraview import caseDirectory as case
from PyFoam.Paraview import viewTime as vTime

from PyFoam.Error import warning

from PyFoam.Applications.SamplePlot import SamplePlot
from PyFoam.Applications.TimelinePlot import TimelinePlot
from PyFoam.Applications.RedoPlot import RedoPlot

from paraview import vtk

from os import path
import re

def checkDir(dName):
    c=case().name
    if path.exists(path.join(c,dName)):
        return dName
    elif path.exists(path.join(c,"postProcessing",dName)):
        return path.join("postProcessing",dName)
    else:
        warning("Could not find",dName,"in",c)
        return dName

def setSampleData(tbl,lineName,dirName="sets"):
    """Read sample data and put it into a vtkTable

    Use in 'Programmable Filter'. Set output type to 'vtkTable'.
    To get (for instance) the data on line 'lineCenter' from the sampled data in
    the directory sets (it will automatically use the current time) in the
    source code write
from PyFoam.Paraview.Data import setSampleData
setSampleData(self.GetOutput(),"lineCenter",dirName="sets")
    """
    actualDir=checkDir(dirName)
    data=SamplePlot(args=[case().name,
                          "--directory="+actualDir,
                          "--time="+str(vTime()),
                          "--line="+lineName,
                          "--fuzzy-time",
                          "--numpy"]).data

    nExp=re.compile(lineName+"_t=[^ ]+ (.+)")

    for n in data.dtype.names:
        m=nExp.match(n)
        vals=data[n]
        if m:
            n=m.group(1)
        col=vtk.vtkDoubleArray()
        col.SetName(n)
        for v in vals:
            col.InsertNextValue(v)
        tbl.AddColumn(col)

    return tbl

def setTimelineData(tbl,dirName,fields=[],positions=[],time=0):
    """Read timeline data and put it into a vtkTable

    Use in 'Programmable Filter'. Set output type to 'vtkTable'.
    To get (for instance) the fields p and U on position 'min' from the timeline data in
    the directory swakExpression_foo in the
    source code write
from PyFoam.Paraview.Data import setTimelineData
setTimelineData(self.GetOutput(),"volumeMinMax_pressureExtreme",positions=["minimum"],fields=["p","magU"])
    """
    actualDir=checkDir(dirName)
    data=TimelinePlot(args=[case().name,
                          "--directory="+actualDir,
                          "--time="+str(time),
                          "--basic-mode=lines",
                          "--numpy"]+
                      ["--field=%s" % f for f in fields]+
                      ["--position=%s" % p for p in positions]).data


    for i,n in enumerate(data.dtype.names):
        vals=data[n]
        col=vtk.vtkDoubleArray()
        col.SetName(n)
        for v in vals:
            col.InsertNextValue(v)
        tbl.AddColumn(col)

    return tbl

def setPlotData(tbl,solverName,plotName):
    """Read plot data and put it into a vtkTable

    Use in 'Programmable Filter'. Set output type to 'vtkTable'.
    To get (for instance) the plot data generated by solver 'compressibleInterFoam'
    (if run with pyFoamRunner.py) under the name 'linear' write
from PyFoam.Paraview.Data import setPlotData
setPlotData(self.GetOutput(),"compressibleInterFoam","linear")
    """
    actualFile=path.join("PyFoamRunner."+solverName+".analyzed","pickledPlots")
    data=RedoPlot(args=[actualFile,
                        "--pickle-file",
                        "--numpy"]).plotNumpy[plotName]

    for n in data.dtype.names:
        vals=data[n]
        col=vtk.vtkDoubleArray()
        col.SetName(n)
        for v in vals:
            col.InsertNextValue(v)
        tbl.AddColumn(col)

    return tbl

def setCurrentTimeline(tbl,dirName,position):
    """Read timeline data at a specified position. Return the interpolated
    values from the current simulation time as a table

    Use in 'Programmable Filter'. Set output type to 'vtkTable'.
    To get the values from a timeline 'postProcessing/swakExpression_pseudoTime/0/pseudoTime' use
from PyFoam.Paraview.Data import setCurrentTimeline
setCurrentTimeline(self.GetOutput(),"swakExpression_pseudoTime","pseudoTime")

    To print the values pipe into a 'PythonAnnotation'-filter. Set
    ArrayAssociation to 'Row Data' and use an Expression like
    '"Time = %.1f (%.1f)" % (average,t_value)' (this assumes that there is a
    column 'average')

    Hint: this reads string values as well. But in that case the value has to be
    converted with 'val.GetValue(0)' in the expression
"""

    from PyFoam.RunDictionary.TimelineDirectory import TimelineDirectory
    from vtk import vtkStringArray,vtkFloatArray

    dirName=checkDir(dirName)
    tld=TimelineDirectory(case().name,dirName)
    spread=tld[position](False)
    data=spread(vTime())
    for k in data:
        try:
            val=float(data[k])
            col=vtkFloatArray()
        except ValueError:
            val=data[k]
            col=vtkStringArray()
        col.SetNumberOfComponents(1)
        col.SetName(k)
        col.InsertNextValue(val)

        tbl.AddColumn(col)

    return col
